#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/3/2 20:47
# @Author  : Seven
# @File    : RTDemo.py
# @Software: PyCharm
# function : 本质矩阵估计两次拍摄间的平移和旋转
import cv2
import numpy as np

# 加载相机标定的数据
with np.load('C.npz') as X:
    mtx, dist, _, _ = [X[i] for i in ('mtx', 'dist', 'rvecs', 'tvecs')]

# 加载图片
img1 = cv2.imread('image/l.jpg', 0)
img2 = cv2.imread('image/r.jpg', 0)

# 初始化SIFT方法
sift = cv2.xfeatures2d_SIFT.create()

# 获取关键点和描述子
k1, d1 = sift.detectAndCompute(img1, None)
k2, d2 = sift.detectAndCompute(img2, None)

# 设置FLANN 超参数
FLANN_INDEX_KDTREE = 0
# K-D树索引超参数
index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
# 搜索超参数
search_params = dict(checks=50)

# 初始化FlannBasedMatcher匹配器
flann = cv2.FlannBasedMatcher(index_params, search_params)
# 通过KNN的方式匹配两张图的描述子
matches = flann.knnMatch(d1, d2, k=2)
# 筛选比较好的匹配点
pts1 = []
pts2 = []
for i, (m, n) in enumerate(matches):
    if m.distance < 0.8 * n.distance:
        pts2.append(k2[m.trainIdx].pt)
        pts1.append(k1[m.queryIdx].pt)

# 计算本质矩阵
pts1 = np.int32(pts1)
pts2 = np.int32(pts2)
# E为本质矩阵、mask是返回基本矩阵的值：没有找到矩阵，返回0，找到一个矩阵返回1，多个矩阵返回3
E, mask = cv2.findEssentialMat(pts1, pts2, mtx)
print("本质矩阵：\n", E)
# 只选择有效数据
pts1 = pts1[mask.ravel() == 1]
pts2 = pts2[mask.ravel() == 1]
# 通过本质矩阵估计两次拍摄间的平移和旋转
ret, R, T, mask = cv2.recoverPose(E, pts1, pts2, mtx)
print("旋转矩阵：\n", R)
print("平移矩阵：\n", T)

# 罗德里格斯旋转矩阵转换
new_R, _ = cv2.Rodrigues(R)
print("Rodrigues转换[R]：\n", new_R)
